from pathlib import Path

from ipyniivue import download_dataset

BASE_API_URL = "https://niivue.com/demos/images"
DATA_FOLDER = Path("images")

# Download data for example
download_dataset(
    BASE_API_URL,
    DATA_FOLDER,
    files=[
        "mni152.nii.gz",
        "tract.FAT_R.vtk",
        "tract.IFOF_R.trk",
        "tract.SLF1_R.tck",
        "BrainMesh_ICBM152.lh.mz3",
    ],
)

mni152.nii.gz already exists.
tract.FAT_R.vtk already exists.
Downloading tract.IFOF_R.trk...
tract.SLF1_R.tck already exists.
BrainMesh_ICBM152.lh.mz3 already exists.
Dataset downloaded successfully to images.

import ipywidgets as widgets
from IPython.display import display

import ipyniivue

# 1. initialize NiiVue
nv = ipyniivue.NiiVue()
nv.opts.show_3d_crosshair = True
nv.opts.back_color = [0, 0, 0, 1]
nv.opts.is_orient_cube = True
nv.opts.slice_type = ipyniivue.SliceType.RENDER

# 2. load volume
nv.add_volume({"path": DATA_FOLDER / "mni152.nii.gz"})

# 3. load meshes
meshes = [
    # Tract 0
    {"path": DATA_FOLDER / "tract.FAT_R.vtk", "rgba255": [180, 180, 0, 255]},
    # Tract 1
    {"path": DATA_FOLDER / "tract.IFOF_R.trk", "rgba255": [0, 255, 0, 255]},
    # Tract 2
    {"path": DATA_FOLDER / "tract.SLF1_R.tck", "rgba255": [0, 0, 255, 255]},
    # Brain Mesh (Index 3)
    {
        "path": DATA_FOLDER / "BrainMesh_ICBM152.lh.mz3",
        "rgba255": [242, 174, 177, 255],
        "opacity": 0.2,
    },
]
nv.load_meshes(meshes)

# 4. set initial scene props
nv.set_clip_plane(-0.1, 0, 90)
nv.set_render_azimuth_elevation(135, 15)

# 5. create ui widgets
shader_drop = widgets.Dropdown(
    options=["Edge", "Flat", "Matcap", "Matte", "Outline", "Phong", "Toon"],
    value="Outline",
    description="Shader:",
)

color_drop = widgets.Dropdown(
    options=[
        ("Global direction", "Global"),
        ("Local direction", "Local"),
        ("Fixed", "Fixed"),
    ],
    value="Fixed",
    description="Color:",
)

fiber_radius = widgets.FloatSlider(
    min=0, max=2.0, step=0.1, value=0, description="Radius"
)

xray_slider = widgets.FloatSlider(
    min=0, max=0.2, step=0.01, value=0.1, description="MeshXRay"
)

zoom_slider = widgets.FloatSlider(
    min=0.1, max=2.0, step=0.1, value=1.0, description="Zoom"
)


# 6. define callbacks
def on_shader_change(change):
    """Apply shader to the brain mesh (index 3)."""
    if len(nv.meshes) > 3:
        nv.set_mesh_shader(nv.meshes[3].id, change["new"])


def on_color_change(change):
    """Apply fiber color to the tracts (indices 0, 1, 2)."""
    val = change["new"]
    for i in range(3):
        if i < len(nv.meshes):
            nv.meshes[i].fiber_color = val


def on_radius_change(change):
    """Apply radius to the tracts (indices 0, 1, 2)."""
    val = change["new"]
    for i in range(3):
        if i < len(nv.meshes):
            nv.meshes[i].fiber_radius = val


def on_xray_change(change):
    """Handle xray change."""
    nv.opts.mesh_xray = change["new"]


def on_zoom_change(change):
    """Handle zoom change."""
    nv.scene.vol_scale_multiplier = change["new"]


# 7. Bind Callbacks
shader_drop.observe(on_shader_change, names="value")
color_drop.observe(on_color_change, names="value")
fiber_radius.observe(on_radius_change, names="value")
xray_slider.observe(on_xray_change, names="value")
zoom_slider.observe(on_zoom_change, names="value")

# 8. Set Initial Values
nv.opts.mesh_xray = xray_slider.value

for i in range(3):
    nv.meshes[i].fiber_color = color_drop.value
    nv.meshes[i].fiber_radius = fiber_radius.value

nv.set_mesh_shader(nv.meshes[3].id, shader_drop.value)


# 9. display all
ui = widgets.VBox(
    [
        widgets.HBox([shader_drop, color_drop]),
        widgets.HBox([fiber_radius, xray_slider, zoom_slider]),
        nv,
    ]
)

display(ui)